"""
训练日志分析脚本
分析JSON格式的训练日志，输出最优模型的详细信息和训练曲线
"""
import os
import sys
import json
import argparse
from datetime import datetime
import matplotlib.pyplot as plt
import numpy as np


def load_log(log_path):
    """
    加载JSON日志文件

    Args:
        log_path: 日志文件路径

    Returns:
        log_data: 日志数据字典
    """
    if not os.path.exists(log_path):
        raise FileNotFoundError(f"日志文件不存在: {log_path}")

    with open(log_path, 'r', encoding='utf-8') as f:
        log_data = json.load(f)

    return log_data


def print_metadata(log_data):
    """打印元数据信息"""
    metadata = log_data.get('metadata', {})

    print("\n" + "=" * 80)
    print("训练元数据")
    print("=" * 80)
    print(f"模型名称: {metadata.get('model_name', 'N/A')}")
    print(f"折数: {metadata.get('fold_idx', 'N/A')}")
    print(f"开始时间: {metadata.get('start_time', 'N/A')}")

    config = metadata.get('config', {})
    if config:
        print("\n训练配置:")
        for key, value in config.items():
            if isinstance(value, dict):
                print(f"  {key}:")
                for k, v in value.items():
                    print(f"    {k}: {v}")
            else:
                print(f"  {key}: {value}")


def print_best_model(log_data):
    """打印最佳模型信息"""
    best_model = log_data.get('best_model', {})

    print("\n" + "=" * 80)
    print("最佳模型信息")
    print("=" * 80)
    print(f"最佳Epoch: {best_model.get('epoch', 'N/A')}")
    print(f"检查点路径: {best_model.get('checkpoint_path', 'N/A')}")
    print(f"保存时间: {best_model.get('timestamp', 'N/A')}")

    metrics = best_model.get('metrics', {})
    if metrics:
        print("\n最佳模型指标:")
        print(f"  {'指标':<20} {'值':<10}")
        print(f"  {'-' * 30}")
        for key, value in sorted(metrics.items()):
            if isinstance(value, (int, float)):
                print(f"  {key:<20} {value:<10.6f}")
            else:
                print(f"  {key:<20} {value}")


def print_summary(log_data):
    """打印训练总结"""
    summary = log_data.get('summary', {})

    print("\n" + "=" * 80)
    print("训练总结")
    print("=" * 80)

    total_epochs = summary.get('total_epochs', 0)
    total_time = summary.get('total_time', 0)
    best_miou = summary.get('best_miou', 0)
    end_time = summary.get('end_time', 'N/A')

    print(f"总Epoch数: {total_epochs}")
    print(f"总训练时间: {total_time:.2f} 秒 ({total_time/3600:.2f} 小时)")
    if total_epochs > 0:
        print(f"平均每Epoch时间: {total_time/total_epochs:.2f} 秒")
    print(f"最佳mIoU: {best_miou:.6f}")
    print(f"结束时间: {end_time}")


def print_training_progress(log_data, last_n=10):
    """
    打印最近N个epoch的训练进度

    Args:
        log_data: 日志数据
        last_n: 显示最近N个epoch
    """
    epochs = log_data.get('training', {}).get('epochs', [])

    if not epochs:
        print("\n[WARNING] 没有训练数据")
        return

    print("\n" + "=" * 80)
    print(f"训练进度（最近{last_n}个Epoch）")
    print("=" * 80)

    # 表头
    print(f"{'Epoch':<8} {'Train Loss':<12} {'Val Loss':<12} {'Val mIoU':<12} {'Val Dice':<12} {'LR':<12} {'Time(s)':<10}")
    print("-" * 80)

    # 显示最近N个epoch
    for epoch_data in epochs[-last_n:]:
        epoch = epoch_data.get('epoch', 0)
        train_loss = epoch_data.get('train', {}).get('loss', 0)
        val_loss = epoch_data.get('val', {}).get('loss', 0)
        val_miou = epoch_data.get('val', {}).get('miou', 0)
        val_dice = epoch_data.get('val', {}).get('dice', 0)
        lr = epoch_data.get('lr', 0)
        time_s = epoch_data.get('time', 0)

        print(f"{epoch:<8} {train_loss:<12.6f} {val_loss:<12.6f} {val_miou:<12.6f} {val_dice:<12.6f} {lr:<12.6e} {time_s:<10.2f}")


def plot_training_curves(log_data, save_path=None):
    """
    绘制训练曲线

    Args:
        log_data: 日志数据
        save_path: 保存路径（可选）
    """
    epochs_data = log_data.get('training', {}).get('epochs', [])

    if not epochs_data:
        print("\n[WARNING] 没有训练数据，无法绘制曲线")
        return

    # 提取数据
    epochs = [e['epoch'] for e in epochs_data]
    train_loss = [e.get('train', {}).get('loss', 0) for e in epochs_data]
    val_loss = [e.get('val', {}).get('loss', 0) for e in epochs_data]
    val_miou = [e.get('val', {}).get('miou', 0) for e in epochs_data]
    val_dice = [e.get('val', {}).get('dice', 0) for e in epochs_data]
    val_iou = [e.get('val', {}).get('iou', 0) for e in epochs_data]

    # 创建图表
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    fig.suptitle(f"训练曲线 - {log_data['metadata']['model_name']} (Fold {log_data['metadata']['fold_idx']})", fontsize=16)

    # Loss曲线
    axes[0, 0].plot(epochs, train_loss, label='Train Loss', marker='o', markersize=3)
    axes[0, 0].plot(epochs, val_loss, label='Val Loss', marker='s', markersize=3)
    axes[0, 0].set_xlabel('Epoch')
    axes[0, 0].set_ylabel('Loss')
    axes[0, 0].set_title('Loss曲线')
    axes[0, 0].legend()
    axes[0, 0].grid(True, alpha=0.3)

    # mIoU曲线
    axes[0, 1].plot(epochs, val_miou, label='Val mIoU', marker='o', markersize=3, color='green')
    axes[0, 1].set_xlabel('Epoch')
    axes[0, 1].set_ylabel('mIoU')
    axes[0, 1].set_title('mIoU曲线')
    axes[0, 1].legend()
    axes[0, 1].grid(True, alpha=0.3)

    # Dice曲线
    axes[1, 0].plot(epochs, val_dice, label='Val Dice', marker='s', markersize=3, color='orange')
    axes[1, 0].set_xlabel('Epoch')
    axes[1, 0].set_ylabel('Dice')
    axes[1, 0].set_title('Dice曲线')
    axes[1, 0].legend()
    axes[1, 0].grid(True, alpha=0.3)

    # IoU曲线
    axes[1, 1].plot(epochs, val_iou, label='Val IoU', marker='^', markersize=3, color='red')
    axes[1, 1].set_xlabel('Epoch')
    axes[1, 1].set_ylabel('IoU')
    axes[1, 1].set_title('IoU曲线')
    axes[1, 1].legend()
    axes[1, 1].grid(True, alpha=0.3)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"\n[INFO] 训练曲线已保存: {save_path}")
    else:
        plt.show()


def export_to_csv(log_data, output_path):
    """
    导出训练数据到CSV

    Args:
        log_data: 日志数据
        output_path: 输出CSV路径
    """
    epochs_data = log_data.get('training', {}).get('epochs', [])

    if not epochs_data:
        print("\n[WARNING] 没有训练数据，无法导出CSV")
        return

    import csv

    with open(output_path, 'w', newline='', encoding='utf-8') as f:
        # 获取所有可能的字段
        fieldnames = ['epoch', 'lr', 'time']

        # 添加训练集字段
        if epochs_data[0].get('train'):
            train_keys = epochs_data[0]['train'].keys()
            fieldnames.extend([f'train_{k}' for k in train_keys])

        # 添加验证集字段
        if epochs_data[0].get('val'):
            val_keys = epochs_data[0]['val'].keys()
            fieldnames.extend([f'val_{k}' for k in val_keys])

        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()

        for epoch_data in epochs_data:
            row = {
                'epoch': epoch_data.get('epoch', 0),
                'lr': epoch_data.get('lr', 0),
                'time': epoch_data.get('time', 0)
            }

            # 添加训练集指标
            for k, v in epoch_data.get('train', {}).items():
                row[f'train_{k}'] = v

            # 添加验证集指标
            for k, v in epoch_data.get('val', {}).items():
                row[f'val_{k}'] = v

            writer.writerow(row)

    print(f"\n[INFO] 训练数据已导出到CSV: {output_path}")


def compare_folds(log_dir, model_name, k_folds=8):
    """
    对比所有fold的结果

    Args:
        log_dir: 日志目录
        model_name: 模型名称
        k_folds: 折数
    """
    print("\n" + "=" * 80)
    print(f"对比所有Fold - {model_name}")
    print("=" * 80)

    results = []

    for fold_idx in range(k_folds):
        log_path = os.path.join(log_dir, model_name, f"{model_name}_fold_{fold_idx}.json")

        if not os.path.exists(log_path):
            print(f"[WARNING] Fold {fold_idx} 日志不存在: {log_path}")
            continue

        log_data = load_log(log_path)
        best_model = log_data.get('best_model', {})
        metrics = best_model.get('metrics', {})

        results.append({
            'fold': fold_idx,
            'epoch': best_model.get('epoch', -1),
            'miou': metrics.get('miou', 0),
            'dice': metrics.get('dice', 0),
            'iou': metrics.get('iou', 0),
            'pa': metrics.get('pixel_accuracy', 0),
            'precision': metrics.get('precision', 0),
            'recall': metrics.get('recall', 0)
        })

    if not results:
        print("[ERROR] 没有找到任何fold的日志")
        return

    # 打印表格
    print(f"\n{'Fold':<6} {'Epoch':<8} {'mIoU':<10} {'Dice':<10} {'IoU':<10} {'PA':<10} {'Precision':<12} {'Recall':<10}")
    print("-" * 80)

    for r in results:
        print(f"{r['fold']:<6} {r['epoch']:<8} {r['miou']:<10.6f} {r['dice']:<10.6f} {r['iou']:<10.6f} "
              f"{r['pa']:<10.6f} {r['precision']:<12.6f} {r['recall']:<10.6f}")

    # 计算统计信息
    if len(results) > 1:
        print("\n" + "-" * 80)
        print("统计信息:")

        for metric in ['miou', 'dice', 'iou', 'pa', 'precision', 'recall']:
            values = [r[metric] for r in results]
            mean = np.mean(values)
            std = np.std(values)
            print(f"  {metric:<12}: {mean:.6f} ± {std:.6f}")


def main():
    parser = argparse.ArgumentParser(description='分析训练日志')
    parser.add_argument('--log', type=str, help='日志文件路径')
    parser.add_argument('--model', type=str, help='模型名称（用于对比所有fold）')
    parser.add_argument('--log_dir', type=str, default='logs', help='日志目录')
    parser.add_argument('--plot', action='store_true', help='绘制训练曲线')
    parser.add_argument('--save_plot', type=str, help='保存训练曲线图片的路径')
    parser.add_argument('--export_csv', type=str, help='导出CSV的路径')
    parser.add_argument('--compare_folds', action='store_true', help='对比所有fold')
    parser.add_argument('--last_n', type=int, default=10, help='显示最近N个epoch')

    args = parser.parse_args()

    # 对比所有fold
    if args.compare_folds and args.model:
        compare_folds(args.log_dir, args.model)
        return

    # 分析单个日志文件
    if not args.log:
        print("[ERROR] 请指定日志文件路径 (--log) 或使用 --compare_folds --model <model_name>")
        return

    try:
        # 加载日志
        log_data = load_log(args.log)

        # 打印信息
        print_metadata(log_data)
        print_best_model(log_data)
        print_summary(log_data)
        print_training_progress(log_data, last_n=args.last_n)

        # 绘制曲线
        if args.plot or args.save_plot:
            plot_training_curves(log_data, save_path=args.save_plot)

        # 导出CSV
        if args.export_csv:
            export_to_csv(log_data, args.export_csv)

        print("\n" + "=" * 80)
        print("分析完成！")
        print("=" * 80 + "\n")

    except Exception as e:
        print(f"[ERROR] 分析失败: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == '__main__':
    main()
